# -*- coding: utf-8 -*-

"""
Step1: AI/LLM 提取文本的核心内容
"""
import os
import sys
from flask import Flask, request, jsonify
import ollama
from transformers import AutoTokenizer, AutoModel
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np
import torch
import json

app = Flask(__name__)

# 输出
def AI_generate(text: str,model):
    stream = ollama.generate(
        stream=True,
        model=model,
        prompt=text,
    )
    response_text = ""
    for chunk in stream:
        if not chunk['done']:
            response_text += chunk['response']
    return response_text

def read_file(file_path: str) -> str:
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read()

def save_to_file(file_path: str, content: str):
    with open(file_path, 'w', encoding='utf-8') as file:
        file.write(content)

# 提取核心内容
def extract_core_content(file_text: str, model: str):
    question = "Please extract the core content of the following text and list the key takeaways: \n" + file_text
    core_content = AI_generate(text=question, model=model)
    return core_content

# 读取并分割长文本
def get_long_text_embedding(text, tokenizer, model, max_length=512, stride=256):
    tokens = tokenizer(text, return_tensors='pt', max_length=max_length, truncation=True, padding='max_length')

    input_ids = tokens['input_ids'][0]
    total_length = len(input_ids)

    all_embeddings = []

    # 使用滑动窗口法，分割长文本并提取嵌入
    for i in range(0, total_length, stride):
        end_index = min(i + max_length, total_length)
        chunk_input_ids = input_ids[i:end_index].unsqueeze(0)

        with torch.no_grad():
            outputs = model(input_ids=chunk_input_ids)

        # 获取 [CLS] token 的嵌入
        cls_embedding = outputs.last_hidden_state[:, 0, :].numpy()
        all_embeddings.append(cls_embedding)

        if end_index == total_length:
            break

    # 平均所有片段的 [CLS] 嵌入
    mean_embedding = np.mean(np.vstack(all_embeddings), axis=0)
    return mean_embedding


# 计算余弦相似度
def calculate_cosine_similarity(embedding1, embedding2):
    return cosine_similarity([embedding1], [embedding2])[0][0]


def cal_B1_B2(file1_path: str, file2_path: str,target):
    try:
        AI_model = "llama3:8b"

        if not os.path.exists(file1_path):
            print(f"file do not exist: {file1_path}")
            return
        if not os.path.exists(file2_path):
            print(f"file do not exist: {file2_path}")
            return

        text1 = read_file(file1_path)
        text2 = read_file(file2_path)

        if target == 'B2':
            # file1 / file2 的 核心内容（AI提取）
            text1 = extract_core_content(text1, AI_model)
            text2 = extract_core_content(text2, AI_model)

        # 计算核心内容的文本相似度
        model_name = "bert-base-uncased"  # 可换
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModel.from_pretrained(model_name)

        embedding1 = get_long_text_embedding(text1, tokenizer, model)
        embedding2 = get_long_text_embedding(text2, tokenizer, model)

        similarity = calculate_cosine_similarity(embedding1, embedding2)
        return float(similarity)

    except Exception as e:
        return str(e), 500


def cal_A1(x):
    if x==1 : return 1.0
    elif x==2 : return 0.5
    elif x>=3 : return 0.0

def cal_A2(x):
    if x==0 or x==1 :
        return 0.0
    elif x>=2 and x<=5:
        return 0.2*(x-1)
    elif x>=6:
        return 1.0

def cal_A3(x,y):
    return x/y

def cal_B3(x):
    return x/100



@app.route('/calCNLead', methods=['POST'])
def calculate_route():
    try:
        data = request.get_json()  # 使用 request.get_json() 方法
        print(data)  # 输出数据用于调试
        if not data:
            return jsonify({"calRes":"No JSON data received"})
        target = data.get('target')
        x = data.get('x')
        y = data.get('y')
        file1_path = data.get('file1_path')
        file2_path = data.get('file2_path')

        if target == 'A1':
            return jsonify({"calRes":cal_A1(int(x))})
        elif target == 'A2':
            return jsonify({"calRes":cal_A2(int(x))})
        elif target == 'A3':
            return jsonify({"calRes":cal_A3(int(x), int(y))})
        elif target == 'B1' or target == 'B2':
            return jsonify({"calRes":cal_B1_B2(file1_path, file2_path, target)})
        elif target == 'B3':
            return jsonify({"calRes":cal_B3(int(x))})
        else:
            return jsonify({"calRes": "Invalid target"})

    except Exception as e:
        return str(e)

if __name__ == '__main__':
    app.run(host='127.0.0.1', port=5000)